import argparse
import numpy as np
from deepsignal3.utils.process_utils import base2code_dna
import pywt

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        default="/home/xiaoyifu/data/plant/rice_ribenqing/train/samples_CHH.hc_poses.signal_bilstm.denoise_fp8.tsv",
        help="trainset input address.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="/home/xiaoyifu/data/plant/rice_ribenqing/train/samples_CHH.hc_poses.signal_bilstm.denoise_fp8.freq.tsv",
        help="trainset output address.",
    )
    return parser.parse_args()


def parse_a_line(line):
    words = line.strip().split()
    k_signals = np.array([[float(y) for y in x.split(",")] for x in words[10].split(";")])
    signal_flantten=k_signals.flatten()
    signal_freqs=np.fft.fft(signal_flantten)
    mag=np.abs(signal_freqs)
    phase=np.angle(signal_freqs)
    wavelet = 'db4'  # 小波基函数，这里选择 Daubechies 4
    level = 4        # 分解级数
    coeffs = pywt.wavedec(signal_flantten, wavelet, level=level)
    # 设计高通滤波器
    high_pass_filter = np.array([1, -1])  # 一阶差分滤波器
    filtered_coeffs = [np.convolve(c, high_pass_filter, mode='same') for c in coeffs]
    reconstructed_signal = pywt.waverec(filtered_coeffs, wavelet)
    features=signal_freqs,mag,phase,reconstructed_signal
    features_str=features_to_str(features)
    row=words+features_str
    return row

def features_to_str(features):
    signal_freqs,mag,phase,reconstructed_signal=features
    signal_freqs_text=';'.join([",".join([str(y) for y in x]) for x in signal_freqs.reshape(21, 16).tolist()])
    mag_text=','.join([str(x) for x in np.around(mag, decimals=6)])
    phase_text=','.join([str(x) for x in np.around(phase, decimals=6)]) 
    reconstructed_signal_text = ';'.join([",".join([str(y) for y in x]) for x in reconstructed_signal.reshape(21, 16).tolist()])
    return [signal_freqs_text,mag_text,phase_text,reconstructed_signal_text]

def parse_data(args):
    fh = open(args.output, "w", buffering=512)
    with open(args.input, "r") as f:
        for line in f:
            row=parse_a_line(line)
            fh.write('\t'.join(row) + '\n')
            
            

if __name__ == '__main__':
    args=parse_args()
    parse_data(args)